"""
26_large_scale_clustering.py
============================
Large-Scale Graph Clustering for 10K+ Nodes

This module covers:
1. Handling graphs with 10,000+ nodes per entity
2. Clustering strategies for scalability
3. Embedding-based partitioning
4. Hierarchical graph organization
5. Memory-efficient operations
6. Performance optimization patterns

Author: Neo4j Tutorial
Date: January 2025
"""

from typing import Dict, List, Any, Optional, Generator
from dataclasses import dataclass
from enum import Enum
import math
import time
from neo4j import Driver
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import get_driver


class ClusteringStrategy(Enum):
    """Strategies for large-scale graph clustering."""
    LOUVAIN = "louvain"
    LABEL_PROPAGATION = "label_prop"
    KMEANS_EMBEDDING = "kmeans_embedding"
    HIERARCHICAL = "hierarchical"
    PARTITIONING = "partitioning"


@dataclass
class ClusterConfig:
    """Configuration for clustering operations."""
    strategy: ClusteringStrategy
    max_cluster_size: int = 1000
    min_cluster_size: int = 10
    target_clusters: Optional[int] = None
    embedding_dimensions: int = 384


@dataclass
class ClusterStats:
    """Statistics for a cluster."""
    cluster_id: int
    size: int
    density: float
    center_node: str
    labels: Dict[str, int]


class LargeScaleGraphManager:
    """
    Manages large-scale graphs with 10,000+ nodes.
    
    When graphs grow large, naive approaches fail:
    - Full graph algorithms become slow O(n²) or worse
    - Memory consumption explodes
    - Query times degrade
    
    Solutions:
    1. CLUSTERING: Group related nodes, process per cluster
    2. HIERARCHICAL: Create summary/supernode layers
    3. PARTITIONING: Split graph into manageable pieces
    4. STREAMING: Process in batches, not all at once
    
    Architecture for Large Patient Graph:
    
    ┌─────────────────────────────────────────────────────────────┐
    │                  HIERARCHICAL GRAPH STRUCTURE               │
    ├─────────────────────────────────────────────────────────────┤
    │                                                             │
    │  Level 3: Region Summaries                                  │
    │  ┌─────────┐    ┌─────────┐    ┌─────────┐                 │
    │  │Northeast│    │ Midwest │    │  West   │                 │
    │  │  5000   │    │  3000   │    │  2000   │                 │
    │  └────┬────┘    └────┬────┘    └────┬────┘                 │
    │       │              │              │                       │
    │  Level 2: Cluster Summaries                                 │
    │  ┌───┴───┐      ┌───┴───┐      ┌───┴───┐                   │
    │  │Cluster│      │Cluster│      │Cluster│                   │
    │  │  500  │ ...  │  300  │ ...  │  200  │                   │
    │  └───┬───┘      └───┬───┘      └───┬───┘                   │
    │      │              │              │                        │
    │  Level 1: Individual Patients                               │
    │  ┌─┬─┬─┬─┐      ┌─┬─┬─┬─┐      ┌─┬─┬─┬─┐                   │
    │  │P│P│P│P│ ...  │P│P│P│P│ ...  │P│P│P│P│                   │
    │  └─┴─┴─┴─┘      └─┴─┴─┴─┘      └─┴─┴─┴─┘                   │
    │                                                             │
    └─────────────────────────────────────────────────────────────┘
    
    Benefits:
    - Query at appropriate level (don't need all 10K patients)
    - Aggregations pre-computed at cluster level
    - Fast navigation: Region -> Cluster -> Patient
    - Memory efficient: Load only needed clusters
    """
    
    def __init__(self, driver: Driver):
        self.driver = driver
    
    # =========================================================
    # CLUSTERING ALGORITHMS
    # =========================================================
    
    def cluster_patients_louvain(
        self,
        write_property: str = "patientCluster"
    ) -> Dict[str, Any]:
        """
        Cluster patients using Louvain algorithm.
        
        Good for finding natural communities based on
        shared diagnoses, medications, doctors, etc.
        
        Args:
            write_property: Property to store cluster ID
            
        Returns:
            Clustering statistics
        """
        # First create a patient similarity graph
        create_similarity_query = """
        // Create patient similarity projection
        CALL gds.graph.project.cypher(
            'patient_similarity',
            'MATCH (p:Patient) RETURN id(p) AS id',
            'MATCH (p1:Patient)-[:DIAGNOSED_WITH]->(d:Disease)<-[:DIAGNOSED_WITH]-(p2:Patient)
             WHERE id(p1) < id(p2)
             WITH p1, p2, count(d) AS sharedDiseases
             WHERE sharedDiseases >= 1
             RETURN id(p1) AS source, id(p2) AS target, sharedDiseases AS weight'
        )
        YIELD graphName, nodeCount, relationshipCount
        RETURN graphName, nodeCount, relationshipCount
        """
        
        louvain_query = """
        CALL gds.louvain.write('patient_similarity', {
            writeProperty: $write_prop,
            relationshipWeightProperty: 'weight'
        })
        YIELD communityCount, modularity, nodePropertiesWritten
        RETURN communityCount, modularity, nodePropertiesWritten
        """
        
        cleanup_query = "CALL gds.graph.drop('patient_similarity', false)"
        
        with self.driver.session() as session:
            try:
                # Drop if exists
                try:
                    session.run(cleanup_query)
                except:
                    pass
                
                # Create projection
                proj_result = session.run(create_similarity_query)
                proj_data = proj_result.single().data()
                
                # Run Louvain
                result = session.run(louvain_query, write_prop=write_property)
                louvain_data = result.single().data()
                
                # Cleanup
                session.run(cleanup_query)
                
                return {
                    "projection": proj_data,
                    "clustering": louvain_data
                }
                
            except Exception as e:
                try:
                    session.run(cleanup_query)
                except:
                    pass
                return {"error": str(e)}
    
    def cluster_by_embeddings(
        self,
        label: str,
        embedding_property: str,
        n_clusters: int = 10
    ) -> Dict[str, Any]:
        """
        Cluster nodes using their embedding vectors.
        
        Uses K-means style clustering in embedding space.
        Good for semantic grouping.
        
        Args:
            label: Node label to cluster
            embedding_property: Property containing embeddings
            n_clusters: Target number of clusters
            
        Returns:
            Clustering results
        """
        # This uses a GDS-based approach or custom implementation
        query = f"""
        // Get embeddings and cluster using GDS k-means (if available)
        // Otherwise, export embeddings for external clustering
        
        MATCH (n:{label})
        WHERE n.{embedding_property} IS NOT NULL
        WITH n, n.{embedding_property} AS embedding
        // Return for external processing or use GDS
        RETURN 
            n.id AS id,
            embedding,
            count(*) OVER () AS total
        LIMIT 10000
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            records = [record.data() for record in result]
            
            # For actual clustering, you would:
            # 1. Export embeddings to numpy
            # 2. Run sklearn KMeans
            # 3. Write cluster IDs back to Neo4j
            
            return {
                "nodes_with_embeddings": len(records),
                "target_clusters": n_clusters,
                "note": "Export embeddings for external clustering or use GDS k-means"
            }
    
    # =========================================================
    # HIERARCHICAL ORGANIZATION
    # =========================================================
    
    def create_cluster_summaries(
        self,
        cluster_property: str = "patientCluster"
    ) -> int:
        """
        Create summary nodes for each cluster.
        
        Summary nodes aggregate cluster statistics:
        - Patient count
        - Common diagnoses
        - Average age
        - Geographic center
        
        Args:
            cluster_property: Property containing cluster IDs
            
        Returns:
            Number of summary nodes created
        """
        query = f"""
        // Create cluster summary nodes
        MATCH (p:Patient)
        WHERE p.{cluster_property} IS NOT NULL
        WITH p.{cluster_property} AS clusterId, collect(p) AS patients
        
        // Calculate aggregate statistics
        WITH clusterId, patients,
             size(patients) AS patientCount,
             avg(p.age) AS avgAge
        
        // Create summary node
        MERGE (cs:ClusterSummary {{clusterId: clusterId}})
        SET cs.patientCount = patientCount,
            cs.avgAge = avgAge,
            cs.updatedAt = datetime()
        
        // Link patients to cluster
        WITH cs, patients
        UNWIND patients AS patient
        MERGE (patient)-[:BELONGS_TO_CLUSTER]->(cs)
        
        RETURN count(DISTINCT cs) AS summariesCreated
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            return result.single()["summariesCreated"]
    
    def create_super_clusters(
        self,
        max_patients_per_super: int = 5000
    ) -> int:
        """
        Create super-clusters grouping multiple clusters.
        
        Adds another hierarchy level for very large graphs.
        
        Args:
            max_patients_per_super: Max patients per super-cluster
            
        Returns:
            Number of super-clusters created
        """
        query = """
        // Group clusters into super-clusters
        MATCH (cs:ClusterSummary)
        WITH cs ORDER BY cs.clusterId
        WITH collect(cs) AS clusters
        
        // Calculate super-cluster assignments
        WITH clusters, 
             size(clusters) AS numClusters,
             $max_size AS maxSize
        
        UNWIND range(0, size(clusters)-1) AS idx
        WITH clusters[idx] AS cluster,
             idx / 10 AS superClusterId  // 10 clusters per super
        
        // Create super-cluster nodes
        MERGE (sc:SuperCluster {id: superClusterId})
        MERGE (cluster)-[:PART_OF]->(sc)
        
        // Aggregate statistics
        WITH sc, sum(cluster.patientCount) AS totalPatients
        SET sc.totalPatients = totalPatients,
            sc.updatedAt = datetime()
        
        RETURN count(DISTINCT sc) AS superClustersCreated
        """
        
        with self.driver.session() as session:
            result = session.run(query, max_size=max_patients_per_super)
            return result.single()["superClustersCreated"]
    
    # =========================================================
    # BATCH PROCESSING
    # =========================================================
    
    def process_in_batches(
        self,
        label: str,
        batch_size: int = 1000,
        process_function: str = None
    ) -> Generator[Dict[str, Any], None, None]:
        """
        Process large node sets in batches.
        
        Prevents memory issues with large graphs.
        
        Args:
            label: Node label to process
            batch_size: Nodes per batch
            process_function: Cypher function to apply
            
        Yields:
            Batch processing results
        """
        # Get total count
        count_query = f"MATCH (n:{label}) RETURN count(n) AS total"
        
        with self.driver.session() as session:
            total = session.run(count_query).single()["total"]
        
        num_batches = math.ceil(total / batch_size)
        
        for batch_num in range(num_batches):
            skip = batch_num * batch_size
            
            batch_query = f"""
            MATCH (n:{label})
            WITH n
            ORDER BY n.id
            SKIP {skip}
            LIMIT {batch_size}
            // Process batch
            WITH collect(n) AS batch
            RETURN size(batch) AS processed, {batch_num} AS batchNum
            """
            
            with self.driver.session() as session:
                result = session.run(batch_query)
                data = result.single().data()
                
                yield {
                    "batch": batch_num + 1,
                    "total_batches": num_batches,
                    "processed": data["processed"],
                    "cumulative": min((batch_num + 1) * batch_size, total)
                }
    
    # =========================================================
    # QUERY OPTIMIZATION FOR LARGE GRAPHS
    # =========================================================
    
    def query_within_cluster(
        self,
        cluster_id: int,
        query_pattern: str
    ) -> List[Dict[str, Any]]:
        """
        Execute query within a specific cluster.
        
        Much faster than querying entire graph.
        
        Args:
            cluster_id: Cluster to query within
            query_pattern: Cypher pattern to match
            
        Returns:
            Query results
        """
        query = f"""
        MATCH (cs:ClusterSummary {{clusterId: $cluster_id}})
        MATCH (p:Patient)-[:BELONGS_TO_CLUSTER]->(cs)
        WITH p
        {query_pattern}
        """
        
        with self.driver.session() as session:
            result = session.run(query, cluster_id=cluster_id)
            return [record.data() for record in result]
    
    def get_cluster_statistics(self) -> List[Dict[str, Any]]:
        """
        Get statistics for all clusters.
        
        Returns:
            List of cluster stats
        """
        query = """
        MATCH (cs:ClusterSummary)
        OPTIONAL MATCH (p:Patient)-[:BELONGS_TO_CLUSTER]->(cs)
        WITH cs, count(p) AS actualCount
        RETURN 
            cs.clusterId AS clusterId,
            cs.patientCount AS storedCount,
            actualCount,
            cs.avgAge AS avgAge,
            cs.updatedAt AS lastUpdate
        ORDER BY actualCount DESC
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            return [record.data() for record in result]
    
    # =========================================================
    # PERFORMANCE MONITORING
    # =========================================================
    
    def benchmark_cluster_query(
        self,
        cluster_id: int,
        iterations: int = 10
    ) -> Dict[str, float]:
        """
        Benchmark query performance within a cluster vs full graph.
        
        Args:
            cluster_id: Cluster to test
            iterations: Number of test iterations
            
        Returns:
            Performance comparison
        """
        # Cluster-scoped query
        cluster_query = """
        MATCH (cs:ClusterSummary {clusterId: $cluster_id})
        MATCH (p:Patient)-[:BELONGS_TO_CLUSTER]->(cs)
        MATCH (p)-[:DIAGNOSED_WITH]->(d:Disease)
        RETURN count(d) AS diagnoses
        """
        
        # Full graph query
        full_query = """
        MATCH (p:Patient)-[:DIAGNOSED_WITH]->(d:Disease)
        RETURN count(d) AS diagnoses
        """
        
        cluster_times = []
        full_times = []
        
        with self.driver.session() as session:
            for _ in range(iterations):
                # Cluster query
                start = time.perf_counter()
                session.run(cluster_query, cluster_id=cluster_id).consume()
                cluster_times.append((time.perf_counter() - start) * 1000)
                
                # Full query
                start = time.perf_counter()
                session.run(full_query).consume()
                full_times.append((time.perf_counter() - start) * 1000)
        
        return {
            "cluster_query_avg_ms": sum(cluster_times) / len(cluster_times),
            "full_query_avg_ms": sum(full_times) / len(full_times),
            "speedup": (sum(full_times) / len(full_times)) / 
                       (sum(cluster_times) / len(cluster_times))
        }


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Large-Scale Graph Clustering - Demonstration")
    print("=" * 60)
    
    print("""
    When handling 10,000+ nodes around a single entity (like Patient):
    
    PROBLEM:
    ─────────
    - Full graph algorithms: O(n²) or worse
    - Memory: Can't load entire graph
    - Queries: Slow traversal times
    
    SOLUTION: HIERARCHICAL CLUSTERING
    ──────────────────────────────────
    
    1. CLUSTER patients by similarity
       - Shared diagnoses
       - Shared medications
       - Geographic proximity
       - Demographic similarity
    
    2. CREATE SUMMARY NODES
       - ClusterSummary: Aggregates for ~500 patients
       - SuperCluster: Groups of clusters
       
    3. QUERY AT RIGHT LEVEL
       - Population stats → SuperCluster level
       - Regional analysis → Cluster level
       - Individual care → Patient level
    
    PERFORMANCE GAINS:
    ─────────────────
    ┌─────────────────────┬──────────────┬──────────────┐
    │ Query Type          │ Full Graph   │ Clustered    │
    ├─────────────────────┼──────────────┼──────────────┤
    │ Count by diagnosis  │ 500ms        │ 50ms         │
    │ Find similar        │ 2000ms       │ 200ms        │
    │ Aggregate stats     │ 1000ms       │ 10ms (cached)│
    └─────────────────────┴──────────────┴──────────────┘
    """)
    
    try:
        driver = get_driver()
        manager = LargeScaleGraphManager(driver)
        
        # Get cluster statistics
        print("\n📌 Cluster Statistics")
        print("-" * 40)
        
        try:
            stats = manager.get_cluster_statistics()
            if stats:
                for s in stats[:5]:
                    print(f"   Cluster {s['clusterId']}: {s['actualCount']} patients")
            else:
                print("   No clusters found. Run clustering first.")
        except Exception as e:
            print(f"   Error getting stats: {e}")
        
        # Batch processing example
        print("\n📌 Batch Processing Example")
        print("-" * 40)
        
        try:
            print("   Processing patients in batches...")
            for batch_info in manager.process_in_batches("Patient", batch_size=100):
                print(f"   Batch {batch_info['batch']}/{batch_info['total_batches']}: "
                      f"{batch_info['processed']} processed")
                if batch_info['batch'] >= 3:  # Stop after 3 batches for demo
                    print("   ... (stopping demo)")
                    break
        except Exception as e:
            print(f"   Error in batch processing: {e}")
        
        driver.close()
        
    except Exception as e:
        print(f"   Error: {e}")
    
    print("\n" + "=" * 60)
    print("✅ Large-scale clustering demonstration complete!")
    print("=" * 60)
